import numpy as np
import g2o


class VertexSE3ExpmapOnlyTrans(g2o.VertexSE3):
    def __init__(self):
        super(VertexSE3ExpmapOnlyTrans, self).__init__()

    def set_to_origin_impl(self):
        self.set_estimate(g2o.SE3Quat())

    def oplus_impl(self, update):
        se3 = self.estimate()
        se3 = se3 * g2o.SE3Quat.exp(update)
        self.set_estimate(se3)


class EdgeSE3ProjectXYZOnlyTrans(g2o.VariableVectorXEdge):
    def __init__(self):
        g2o.VariableVectorXEdge.__init__(self)
        self.set_dimension(2)  # dimension of the error function
        self.information()
        self.resize(1)  # number of vertices
        self.set_measurement([0, 0, 0])  # initial measurement

    def linearize_oplus(self):
        vj = self.vertex(1)
        T = vj.estimate()
        vi = self.vertex(0)
        xyz = vi.estimate()
        xyz_trans = T.map(xyz)

        x = xyz_trans[0]
        y = xyz_trans[1]
        z = xyz_trans[2]
        z_2 = z * z

        cam = self.parameter(0)

        tmp = np.zeros((2, 3))
        tmp[0, 0] = cam.focal_length
        tmp[0, 1] = 0
        tmp[0, 2] = -x / z * cam.focal_length

        tmp[1, 0] = 0
        tmp[1, 1] = cam.focal_length
        tmp[1, 2] = -y / z * cam.focal_length

        self.jacobianOplusXi = -1.0 / z * tmp @ T.rotation().matrix()
        self.jacobianOplusXi.fill(0)

        self.jacobianOplusXj[0, 0] = x * y / z_2 * cam.focal_length
        self.jacobianOplusXj[0, 1] = -(1 + (x * x / z_2)) * cam.focal_length
        self.jacobianOplusXj[0, 2] = y / z * cam.focal_length
        self.jacobianOplusXj[0, 3] = -1.0 / z * cam.focal_length
        self.jacobianOplusXj[0, 4] = 0
        self.jacobianOplusXj[0, 5] = x / z_2 * cam.focal_length

        self.jacobianOplusXj[1, 0] = (1 + y * y / z_2) * cam.focal_length
        self.jacobianOplusXj[1, 1] = -x * y / z_2 * cam.focal_length
        self.jacobianOplusXj[1, 2] = -x / z * cam.focal_length
        self.jacobianOplusXj[1, 3] = 0
        self.jacobianOplusXj[1, 4] = -1.0 / z * cam.focal_length
        self.jacobianOplusXj[1, 5] = y / z_2 * cam.focal_length
        self.jacobianOplusXj[:2, :3].fill(0)
